//
//  YuanXinURLProtocol.m
//  yuanxinwebview
//
//  Created by GJS on 2019/1/2.
//  Copyright © 2019 晏德智. All rights reserved.
//

#import "YuanXinURLProtocol.h"
#import "JWCacheURLProtocol.h"
#import "CPURLSessionChallengeSender.h"
#import "NSURLRequest+CYLNSURLProtocolExtension.h"

static NSString * const URLProtocolAlreadyHandleKey = @"alreadyHandle";

@interface YuanXinURLProtocol ()

@property (readwrite, nonatomic, strong) NSOperationQueue *forgeroundNetQueue;
@property (readwrite, nonatomic, strong) NSOperationQueue *backgroundNetQueue;

@property (readwrite, nonatomic, strong) NSURLSession *session;
@property (readwrite, nonatomic, strong) NSMutableData *data;
@property (readwrite, nonatomic, strong) NSURLResponse *response;

@property (readwrite, nonatomic, assign) BOOL useCache;
@property (readwrite, nonatomic, copy) NSString *userName;
@property (readwrite, nonatomic, copy) NSString *password;

@property (nullable, nonatomic, strong) NSMutableDictionary<NSString *, NSString *> *postRequestBodies;//缓存post请求body

@end

@implementation YuanXinURLProtocol

+ (instancetype)instance{
    static YuanXinURLProtocol *urlProtocol = nil;
    static dispatch_once_t onceToken;
    dispatch_once(&onceToken, ^{
        urlProtocol = [[YuanXinURLProtocol alloc] init];
        urlProtocol.useCache = NO;
    });
    return urlProtocol;
}

// ---------------- Super ----------------
+ (BOOL)registerClass:(Class)protocolClass {
    return [super registerClass:protocolClass];
}

+ (void)unregisterClass:(Class)protocolClass {
    return [super unregisterClass:protocolClass];
}

+ (BOOL)canInitWithRequest:(NSURLRequest *)request{
    NSString *urlScheme = [[request URL] scheme];
    if ([urlScheme caseInsensitiveCompare:@"http"] == NSOrderedSame || [urlScheme caseInsensitiveCompare:@"https"] == NSOrderedSame){
        //看看是否已经处理过了，防止无限循环
        if ([NSURLProtocol propertyForKey:URLProtocolAlreadyHandleKey inRequest:request]) {
            return NO;
        }
        if ([[self instance] useCache]) { // 缓存功能换用 JWCacheURLProtocol
            return NO;
        }
        return YES;
    }
    return NO;
}

+ (NSURLRequest *)canonicalRequestForRequest:(NSURLRequest *)request{
    //return request;
    NSMutableURLRequest *mRequest = [[request cyl_getPostRequestIncludeBody] mutableCopy];
    if ([request.HTTPMethod isEqualToString:@"POST"] && (mRequest.HTTPBody.length == 0)) {
        NSMutableDictionary *bodyDict = [[self instance] postRequestBodies];
        if (bodyDict) {
            // 有风险。。
            // 而且很可能是相对路径。。
            for (NSString *url in bodyDict.allKeys) {
                NSString *path = [[url componentsSeparatedByString:@"../"] lastObject];
                path = [path stringByReplacingOccurrencesOfString:@"./" withString:@""];
                if ([mRequest.URL.absoluteString containsString:path]) {
                    NSString *bodyData = [bodyDict objectForKey:url];
                    //[bodyDict removeObjectForKey:url]; // 不用移除，token可能会一直有请求。。
                    if ([bodyData isKindOfClass:NSData.class]) {
                        mRequest.HTTPBody = (NSData *)bodyData;
                    } else if ([bodyData isKindOfClass:NSString.class]) {
                        mRequest.HTTPBody = [bodyData dataUsingEncoding:NSUTF8StringEncoding];
                    }
                }
            }
        }
    }
    return [mRequest copy];
}

+ (BOOL)requestIsCacheEquivalent:(NSURLRequest *)a toRequest:(NSURLRequest *)b {
    return [super requestIsCacheEquivalent:a toRequest:b];
}

+ (nullable id)propertyForKey:(NSString *)key inRequest:(NSURLRequest *)request {
    return [super propertyForKey:key inRequest:request];
}

+ (void)setProperty:(id)value forKey:(NSString *)key inRequest:(NSMutableURLRequest *)request {
    return [super setProperty:value forKey:key inRequest:request];
}

+ (void)removePropertyForKey:(NSString *)key inRequest:(NSMutableURLRequest *)request {
    return [super removePropertyForKey:key inRequest:request];
}

// ---------------- Custom ----------------
+ (void)startListeningNetWorking{
    [NSURLProtocol registerClass:[YuanXinURLProtocol class]];
}

+ (void)cancelListeningNetWorking{
    [NSURLProtocol unregisterClass:[YuanXinURLProtocol class]];
    [JWCacheURLProtocol cancelListeningNetWorking];
    [[self instance] setRequestHeaders:nil];
    [[self instance] setPostRequestBodies:nil];
}

+ (void)setConfig:(NSURLSessionConfiguration *)config{
    [[self instance] setConfig:config];
    [JWCacheURLProtocol setConfig:config];
}

+ (void)setRequestHeaders:(NSDictionary<NSString *, NSString *> *)requestHeaders{
    // NSURLSessionConfiguration 有个 HTTPAdditionalHeaders，没有试过这个东西
    [[self instance] setRequestHeaders:requestHeaders];
    [JWCacheURLProtocol setRequestHeaders:requestHeaders];
}

+ (void)setUserName:(NSString *)userName {
    [[self instance] setUserName:userName];
    [JWCacheURLProtocol setUserName:userName];
}

+ (void)setPassword:(NSString *)password {
    [[self instance] setPassword:password];
    [JWCacheURLProtocol setPassword:password];
}

+ (void)setUseCache:(BOOL)ifUseCache { // 需要缓存功能时切换 JWCacheURLProtocol
    YuanXinURLProtocol *instance = [self instance];
    instance.useCache = ifUseCache;
    NSDictionary *requestHeaders = [instance.requestHeaders copy]; // 先保存 requestHeaders，因为 cancelListeningNetWorking 里面写了清空
    if (instance.useCache) {
        [self cancelListeningNetWorking];
        [JWCacheURLProtocol startListeningNetWorking];
        [JWCacheURLProtocol setRequestHeaders:requestHeaders];
    } else {
        [JWCacheURLProtocol cancelListeningNetWorking];
        [self startListeningNetWorking];
    }
}

+ (void)setPostRequestBodies:(NSDictionary<NSString *,NSString *> *)postRequestBodies {
    NSMutableDictionary *mDict = [[self instance] postRequestBodies];
    [mDict addEntriesFromDictionary:postRequestBodies];
    [JWCacheURLProtocol setPostRequestBodies:postRequestBodies];
}

+ (BOOL)appendHeaders:(NSDictionary<NSString *, NSString *> *)headers toRequest:(NSMutableURLRequest *)request {
    BOOL merged = NO;
    for (NSString *key in headers) {
        if (headers[key]) {
            if (![[request.allHTTPHeaderFields allKeys] containsObject:key]) {
                merged = YES;
                [request setValue:headers[key] forHTTPHeaderField:key];
            }
        }
    }
    
    return merged;
}

+ (BOOL)mergeHeaders:(NSDictionary<NSString *, NSString *> *)headers toRequest:(NSMutableURLRequest *)request {
    BOOL merged = NO;
    for (NSString *key in headers) {
        if (headers[key]) {
            if (![request.allHTTPHeaderFields[key] isEqualToString:headers[key]]) {
                merged = YES;
            }
            [request setValue:headers[key] forHTTPHeaderField:key];
        }
    }
    
    return merged;
}

- (NSMutableDictionary<NSString *,NSString *> *)postRequestBodies {
    if (_postRequestBodies == nil) {
        _postRequestBodies = [NSMutableDictionary new];
    }
    
    return _postRequestBodies;
}

- (NSURLSessionConfiguration *)config{
    if (!_config) {
        NSURLSessionConfiguration *configuration = [NSURLSessionConfiguration defaultSessionConfiguration];
        configuration.timeoutIntervalForRequest = 60;
        configuration.HTTPCookieAcceptPolicy = NSHTTPCookieAcceptPolicyAlways;
        NSArray *protocolArray = @[ [YuanXinURLProtocol class] ];
        configuration.protocolClasses = protocolArray;
        
        /* To Fixed AIRWatch SDK Issue in API Calling */
        
        // Create an NSURLSessionConfiguration that uses the proxy
        
//        NSDictionary *proxyDict = @{
//
//                                    @"HTTPEnable"  : [NSNumber numberWithInt:1],
//
//                                    @"HTTPSEnable" : [NSNumber numberWithInt:1],
//
//                                    };
//
//        configuration.connectionProxyDictionary = proxyDict;
        
        _config = configuration;
    }
    return _config;
}

- (NSOperationQueue *)forgeroundNetQueue{
    if (!_forgeroundNetQueue) {
        _forgeroundNetQueue = [[NSOperationQueue alloc] init];
        _forgeroundNetQueue.maxConcurrentOperationCount = 10;
    }
    return _forgeroundNetQueue;
}

- (NSOperationQueue *)backgroundNetQueue{
    if (!_backgroundNetQueue) {
        _backgroundNetQueue = [[NSOperationQueue alloc] init];
        _backgroundNetQueue.maxConcurrentOperationCount = 6;
    }
    return _backgroundNetQueue;
}

- (void)netRequestWithRequest:(NSURLRequest *)request{
    NSURLSessionConfiguration *config = self.config;
    //self.session = [NSURLSession sessionWithConfiguration:config delegate:self delegateQueue:self.forgeroundNetQueue];
    self.session = [NSURLSession sessionWithConfiguration:config delegate:self delegateQueue:nil];
    NSURLSessionDataTask * sessionTask = [self.session dataTaskWithRequest:request];
    [sessionTask resume];
}

- (void)startLoading{
    
    NSMutableURLRequest *mutableRequest = [[self request] mutableCopy];
    
    // 添加自定义请求头
    [self.class appendHeaders:[[self.class instance] requestHeaders] toRequest:mutableRequest];
    
    // 标示该request已经处理过了，防止无限循环
    [NSURLProtocol setProperty:@YES forKey:URLProtocolAlreadyHandleKey inRequest:mutableRequest];
    
    [self netRequestWithRequest:mutableRequest];
}

- (void)stopLoading{
    [self.session invalidateAndCancel];
    self.session = nil;
}

- (void)appendData:(NSData *)newData
{
    if ([self data] == nil) {
        [self setData:[newData mutableCopy]];
    }
    else {
        [[self data] appendData:newData];
    }
}
#pragma mark -NSURLSessionDataDelegate

- (void)URLSession:(NSURLSession *)session dataTask:(NSURLSessionDataTask *)dataTask didReceiveData:(NSData *)data{
    
    [self.client URLProtocol:self didLoadData:data];
    
    [self appendData:data];
}

- (void)URLSession:(NSURLSession *)session dataTask:(NSURLSessionDataTask *)dataTask didReceiveResponse:(NSURLResponse *)response completionHandler:(void (^)(NSURLSessionResponseDisposition))completionHandler{
    [self.client URLProtocol:self didReceiveResponse:response cacheStoragePolicy:NSURLCacheStorageNotAllowed];
    self.response = response;
    completionHandler(NSURLSessionResponseAllow);
}

- (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task didCompleteWithError:(NSError *)error{
    if (error) {
        [self.client URLProtocol:self didFailWithError:error];
    } else {
        [self.client URLProtocolDidFinishLoading:self];
        if (!self.data) {
            return;
        }
        NSCachedURLResponse *cacheUrlResponse = [[NSCachedURLResponse alloc] initWithResponse:task.response data:self.data];
        [[NSURLCache sharedURLCache] storeCachedResponse:cacheUrlResponse forRequest:self.request];
        self.data = nil;
    }
}

//重定向
- (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task willPerformHTTPRedirection:(NSHTTPURLResponse *)response newRequest:(NSURLRequest *)request completionHandler:(void (^)(NSURLRequest * _Nullable))completionHandler {
    
    //[self.client URLProtocol:self wasRedirectedToRequest:request redirectResponse:response];
    
    NSMutableURLRequest* redirectRequest;
    redirectRequest = [request mutableCopy];
    [[self class] removePropertyForKey:URLProtocolAlreadyHandleKey inRequest:redirectRequest];
    [[self client] URLProtocol:self wasRedirectedToRequest:redirectRequest redirectResponse:response];
    
    [task cancel];
    [[self client] URLProtocol:self didFailWithError:[NSError errorWithDomain:NSCocoaErrorDomain code:NSUserCancelledError userInfo:nil]];
}

- (void)URLSession:(NSURLSession *)session dataTask:(NSURLSessionDataTask *)dataTask willCacheResponse:(NSCachedURLResponse *)proposedResponse completionHandler:(void (^)(NSCachedURLResponse * _Nullable))completionHandler {
    completionHandler(proposedResponse);
}

- (void)URLSession:(NSURLSession *)session didReceiveChallenge:(NSURLAuthenticationChallenge *)challenge completionHandler:(void (^)(NSURLSessionAuthChallengeDisposition, NSURLCredential *))completionHandler{
    
    YuanXinURLProtocol *instance = [self.class instance];
    // 判断服务器采用的验证方法
    if ([challenge.protectionSpace.authenticationMethod isEqualToString:NSURLAuthenticationMethodNTLM])
    {
        if ([challenge previousFailureCount] == 0) {
            /* SET YOUR credentials, i'm just hard coding them in, tweak as necessary */
            if (instance.userName && instance.password) {
                NSURLCredential *credential = [NSURLCredential credentialWithUser:instance.userName password:instance.password persistence:NSURLCredentialPersistencePermanent];
                completionHandler(NSURLSessionAuthChallengeUseCredential,credential);
            } else {
                completionHandler(NSURLSessionAuthChallengePerformDefaultHandling, nil);
            }
        } else {
            completionHandler(NSURLSessionAuthChallengeCancelAuthenticationChallenge, nil);
        }
    } else if (challenge.protectionSpace.authenticationMethod == NSURLAuthenticationMethodServerTrust) {
        // 如果没有错误的情况下 创建一个凭证，并使用证书
        if (challenge.previousFailureCount == 0) {
            NSURLCredential *credential = [NSURLCredential credentialForTrust:challenge.protectionSpace.serverTrust];
            completionHandler(NSURLSessionAuthChallengeUseCredential, credential);
        }else {
            // 验证失败，取消本次验证
            completionHandler(NSURLSessionAuthChallengeCancelAuthenticationChallenge, nil);
        }
    }else {
        completionHandler(NSURLSessionAuthChallengeCancelAuthenticationChallenge, nil);
    }
    
    NSURLAuthenticationChallenge* challengeWrapper = [[NSURLAuthenticationChallenge alloc] initWithAuthenticationChallenge:challenge sender:[[CPURLSessionChallengeSender alloc] initWithSessionCompletionHandler:completionHandler]];
    [self.client URLProtocol:self didReceiveAuthenticationChallenge:challengeWrapper];
}

- (NSURLCredential *)handleChallengeServerTrust:(NSURLAuthenticationChallenge *)challenge {
    /*
     //直接验证服务器是否被认证（serverTrust），这种方式直接忽略证书验证，直接建立连接，但不能过滤其它URL连接，可以理解为一种折衷的处理方式，实际上并不安全，因此不推荐。
     SecTrustRef serverTrust = [[challenge protectionSpace] serverTrust];
     return [[challenge sender] useCredential: [NSURLCredential credentialForTrust: serverTrust]
     forAuthenticationChallenge: challenge];
     */
    if ([[[challenge protectionSpace] authenticationMethod] isEqualToString: NSURLAuthenticationMethodServerTrust]) {
        do
        {
            SecTrustRef serverTrust = [[challenge protectionSpace] serverTrust];
            //NSCAssert(serverTrust != nil, @"serverTrust is nil");
            if(nil == serverTrust)
                break; /* failed */
            /**
             *  导入多张CA证书（Certification Authority，支持SSL证书以及自签名的CA）
             */
            NSString *cerPath = [[NSBundle mainBundle] pathForResource:@"client" ofType:@"cer"];//自签名证书
            NSData* caCert = [NSData dataWithContentsOfFile:cerPath];
            
            NSString *cerPath2 = [[NSBundle mainBundle] pathForResource:@"client1" ofType:@"cer"];//自签名证书
            NSData * caCert2 = [NSData dataWithContentsOfFile:cerPath2];
            
            //NSCAssert(caCert != nil, @"caCert is nil");
            if(nil == caCert)
                break; /* failed */
            
            SecCertificateRef caRef = SecCertificateCreateWithData(NULL, (__bridge CFDataRef)caCert);
            //NSCAssert(caRef != nil, @"caRef is nil");
            if(nil == caRef)
                break; /* failed */
            
            SecCertificateRef caRef2 = SecCertificateCreateWithData(NULL, (__bridge CFDataRef)caCert2);
            //NSCAssert(caRef2 != nil, @"caRef2 is nil");
            
            NSArray *caArray = @[(__bridge id)(caRef)];
            if(nil != caRef2) {
                caArray = @[(__bridge id)(caRef),(__bridge id)(caRef2)];
            }
            
            //NSCAssert(caArray != nil, @"caArray is nil");
            if(nil == caArray)
                break; /* failed */
            
            OSStatus status = SecTrustSetAnchorCertificates(serverTrust, (__bridge CFArrayRef)caArray);
            //NSCAssert(errSecSuccess == status, @"SecTrustSetAnchorCertificates failed");
            if(!(errSecSuccess == status))
                break; /* failed */
            
            SecTrustResultType result = -1;
            status = SecTrustEvaluate(serverTrust, &result);
            if(!(errSecSuccess == status))
                break; /* failed */
            NSLog(@"stutas:%d",(int)status);
            NSLog(@"Result: %d", result);
            
            BOOL allowConnect = (result == kSecTrustResultUnspecified) || (result == kSecTrustResultProceed);
            if (allowConnect) {
                NSLog(@"success");
            }else {
                NSLog(@"error");
            }
            /* https://developer.apple.com/library/ios/technotes/tn2232/_index.html */
            /* https://developer.apple.com/library/mac/qa/qa1360/_index.html */
            /* kSecTrustResultUnspecified and kSecTrustResultProceed are success */
            if(! allowConnect)
            {
                break; /* failed */
            }
            
#if 0
            /* Treat kSecTrustResultConfirm and kSecTrustResultRecoverableTrustFailure as success */
            /*   since the user will likely tap-through to see the dancing bunnies */
            if(result == kSecTrustResultDeny || result == kSecTrustResultFatalTrustFailure || result == kSecTrustResultOtherError)
                break; /* failed to trust cert (good in this case) */
#endif
            
            // The only good exit point
            //return [[challenge sender] useCredential: [NSURLCredential credentialForTrust: serverTrust] forAuthenticationChallenge: challenge];
            return [NSURLCredential credentialForTrust: serverTrust];
            
        } while(0);
    }
    
    return nil;
}

@end
